import os.path

import pandas as pd
import numpy as np
from datetime import date
from datetime import timedelta
from pandas import DatetimeIndex
from pandas.tseries.offsets import MonthEnd
from pandas import to_datetime
from tqdm import tqdm
from utils import save_pickle


class DateTransformer:
    """
    实例化日期类,用于生成所需要定位的日期,默认是今日
    td 今日
    mfd 本月第一天
    mld 本月最后一天
    lmfd 上个月第一天
    lmld 上个月最后一天
    nmfd 下个月第一天
    nmld 下个月最后一天
    lymfd 上年同月第一天
    lymld 上年同月最后一天

    """

    def __init__(self, td=date.today()):
        super(DateTransformer, self).__init__()
        # 本日,本月第一天,本月最后一天,上年本月第一天,上年本月最后一天
        self.today = to_datetime(td)
        self.month_first_day = date(self.today.year, self.today.month, 1)
        self.month_last_day = self.today + MonthEnd(0)
        self.last_month_last_day = self.month_first_day + timedelta(days=-1)
        self.last_month_first_day = date(self.last_month_last_day.year, self.last_month_last_day.month, 1)
        self.next_month_first_day = self.get_next_month_first_day()
        self.next_month_last_day = self.next_month_first_day + MonthEnd(0)
        self.last_year_month_first_day = date(self.today.year - 1, self.today.month, 1)
        self.last_year_month_last_day = self.last_year_month_first_day + MonthEnd(0)
        self.dt = {'td': self.today.strftime("%Y-%m-%d"),
                   'mfd': self.month_first_day.strftime("%Y-%m-%d"),
                   'mld': self.month_last_day.strftime("%Y-%m-%d"),
                   'lmfd': self.last_month_first_day.strftime("%Y-%m-%d"),
                   'lmld': self.last_month_last_day.strftime("%Y-%m-%d"),
                   'nmfd': self.next_month_first_day.strftime("%Y-%m-%d"),
                   'nmld': self.next_month_last_day.strftime("%Y-%m-%d"),
                   'lymfd': self.last_year_month_first_day.strftime("%Y-%m-%d"),
                   'lymld': self.last_year_month_last_day.strftime("%Y-%m-%d")}

    def get_next_month_first_day(self):
        """
        获取下个月的第一天（有可能会遇到跨年，所以逻辑复杂了一些）
        当前实现方式不够精简，待优化
        :return:
        """
        if self.today.month == 12:
            month = 1
            year = self.today.year + 1
        else:
            month = self.today.month + 1
            year = self.today.year
        return date(year, month, 1)


class DataTransformer:
    """
    数据频率有 日 周 旬 月 季 年
    数据类型有 原始值 同比 环比
    数据种类 量 价 比率
    输入的字典类型，案例如下：
    {
    '__id': 'D2455089',
     'df': 'dataframe',
     'name': '美国:TSA安检人数',
     '国家': '美国',
     '指标ID': 'D2455089',
     '时间区间': '20210101:20220220',
     '指标名称': '美国:TSA安检人数',
     '来源': '美国交通安全管理局',
     '频率': '日',
     '更新时间': '2022-02-21',
     '单位': '人',
     'x': 1,
     'y': 0
     }
    """

    def __init__(self, data_dic: dict):
        super(DataTransformer, self).__init__()

        self.freq_li_en = ['D', 'W', 'U', 'M', 'Q', 'Y']
        self.freq_li = ['日', '周', '旬', '月', '季', '年']
        self.name_map = {'平均': 'mean', '最后': 'last'}
        self.freq_map = dict(zip(self.freq_li, self.freq_li_en))
        self.data_dic = data_dic
        for k, v in data_dic.items():
            self.__dict__[k] = v

    def auto_detect_freq(self) -> str:
        """
        通过数据间的间隔自动监测数据频率（按照经验参数）
        :return:str
        """
        idx = sorted(self.data_dic['df'].index)
        if not isinstance(idx, DatetimeIndex):
            idx = DatetimeIndex(idx)
        mean_delta = np.mean([(idx[i] - idx[i - 1]).days for i in range(1, len(idx) - 1)])
        if mean_delta > 350:
            return '年'
        elif mean_delta > 85:
            return '季'
        elif mean_delta > 29:
            return '月'
        elif mean_delta > 9:
            return '旬'
        elif mean_delta > 5:
            return '周'
        else:
            return '日'

    @staticmethod
    def get_dic_by_df(df):
        cp = df.copy()
        cp['month_year'] = [f"{s.year}_{s.month}" for s in df.index]
        dic = {}
        for month_year, df_daily in cp.groupby(['month_year']):
            dic[str(df_daily.index[-1])[:10]] = df_daily.iloc[:, :1]
        return dic

    def __getitem__(self, item):
        """
        :param item:月_平均_同比_lag_12（表示求日度数据平均值的当月同比，并滞后12个月）
                    月_平均_值_lag_2 , 月_最后_值_lag_0
                    月_平均_环比_lag_0 , 月_最后_环比_lag_12
        :return:
        """
        name_save = self.name.replace(':', '_')
        pkl_path = f"x_cache_data/{name_save}_{item}.pkl"
        if os.path.exists(pkl_path):
            return pd.read_pickle(pkl_path)
        else:
            ags = item.split('_')
            mean_last = self.name_map[ags[1]]
            lag_ma = ags[3]
            lag = int(ags[-1])
            if ags[2] == '值':
                df = self.to_freq('日', mean_last)
                dic = self.get_dic_by_df(df)
            elif ags[2] == '同比':
                dic, df = self.daily_yoy(mean_last)
            elif ags[2] == '环比':
                dic, df = self.daily_yoy(mean_last, mom=True)
            else:
                print('索引错误')
            df.index = pd.to_datetime(df.index)
            if lag == 0:
                save_pickle((dic, df), pkl_path)
                return dic, df
            else:
                lag_df = df.shift(lag, 'M')
                lag_dic = {}
                for t in lag_df.index:
                    se = lag_df.loc[t, :]
                    t_li = pd.date_range(DateTransformer(t).month_first_day, t, freq='D')
                    concater = []
                    for t_daily in t_li:
                        new_se = se.copy()
                        new_se.name = t_daily
                        concater.append(new_se)
                    daily_df = pd.concat(concater, axis=1).T
                    lag_dic[str(t)[:10]] = daily_df
                save_pickle((lag_dic, lag_df), pkl_path)
                return lag_dic, lag_df

    @property
    def freq(self) -> str:
        """
        获取数据频率属性
        :return:
        """
        if '频率' in self.data_dic:
            freq = self.data_dic['频率']
            if freq not in self.freq_li:
                freq = self.auto_detect_freq()
        else:
            freq = self.auto_detect_freq()
        return freq

    @property
    def mode(self) -> str:
        """
        获取数据类型属性
        :return:
        """
        if '收益率' in self.data_dic['name']:
            return '比率'
        elif '指数' in self.data_dic['name']:
            return '价'
        else:
            return '量'

    @property
    def unit(self) -> str:
        """
        获取数据单位属性
        :return:
        """
        if '单位' in self.data_dic:
            return self.data_dic['单位']
        else:
            return '点'

    def high_freq_to_low(self, df: pd.DataFrame, freq: str, how: str = 'last') -> pd.DataFrame:
        """
        高频数据向低频数据转换
        :param df: DataFrame格式的数据
        :param freq: str 日 周 等
        :param how: last mean last表示数据最后一点， mean表示累计平均
        :return: DataFrame 新的频率的数据
        """
        idx = df.index.to_list()
        t_end = DateTransformer(idx[-1]).dt['mld']
        new_index = pd.date_range(start=idx[0], end=t_end, freq=freq)
        new_df = pd.DataFrame(index=new_index, columns=df.columns)
        for i, t in enumerate(new_index):
            if i == 0:
                slice_df = df.loc[:t, :]
            else:
                slice_df = df.loc[new_index[i - 1]:t, :]

            fill_df = slice_df.fillna(method='ffill')

            if len(fill_df) == 0:
                new_df.loc[t, :] = np.nan
                print('warning', self.data_dic['name'], f'本{freq}无数据>@t={t}')
            else:
                if how == 'last':
                    slice_se = fill_df.iloc[-1, :]
                else:
                    slice_se = fill_df.mean(axis=0)
                new_df.loc[t, :] = slice_se.copy()
        new_df.fillna(method='ffill', inplace=True)
        return new_df

    @staticmethod
    def low_freq_to_high(df: pd.DataFrame, freq: str, method: str = 'ffill') -> pd.DataFrame:
        """
        低频数据向高频数据转换
        :param df: DataFrame格式的数
        :param freq: str 月 季
        :param method: ffill bfill 数据填充方式，一般只能向前填充
        :return: DataFrame
        """
        idx = df.index.to_list()
        new_index = pd.date_range(start=idx[0], end=idx[-1], freq=freq)
        new_df = pd.DataFrame(index=new_index, columns=df.columns)
        for t in idx:
            new_df.loc[t, :] = df.loc[t, :]
        if method == 'ffill':
            new_df.fillna(method='ffill', inplace=True)
        else:
            new_df.interpolate(method='linear', inplace=True)
        return new_df

    def to_freq(self, target_freq: str, how='last', method='ffill') -> pd.DataFrame:
        """
        变频函数
        :param target_freq: 频率[日，周，月，季,年]
        :param how: 【降频用参数】高频到低频取值 last：最后一个值 mean：平均值
        :param method: 【提频用参数】ffill前值填充，interpolate插值填充（插值会泄露信息）
        :return: DataFrame
        """
        assert target_freq not in ['旬']
        raw_freq_idx = self.freq_li.index(self.freq)
        target_freq_idx = self.freq_li.index(target_freq)
        if raw_freq_idx > target_freq_idx:
            return self.low_freq_to_high(df=self.data_dic['df'], freq=self.freq_map[target_freq], method=method)
        elif raw_freq_idx == target_freq_idx:
            df = self.data_dic['df']
            return df.fillna(method='ffill').fillna(method='bfill')
        else:
            return self.high_freq_to_low(df=self.data_dic['df'], freq=self.freq_map[target_freq], how=how)

    def cum_to_mon(self):
        """
        累计值到当月同比的数据转换，一般没有一月数据
        :return:
        """
        df = self.data_dic['df']
        df_mon = df.copy().iloc[1:, :]
        for t in df_mon.index:
            t = pd.to_datetime(t)
            date_tran = DateTransformer(t)
            if t.month > 2:
                t_last_month = pd.to_datetime(date_tran.last_month_last_day)
                df_mon.loc[t, :] = df.loc[t, :] - df.loc[t_last_month, :]
        df_yoy = df_mon.copy().iloc[12:, :]
        for t in df_yoy.index:
            t_last_year = pd.to_datetime(DateTransformer(t).last_year_month_last_day)
            df_yoy.loc[t, :] = df_mon.loc[t, :] / df_mon.loc[t_last_year, :] - 1
        return df_yoy

    def to_yoy(self, freq):
        """
        数据转换为月度同比；
        高频流量数据通过求和变频到月度，然后求同比
        高频的指数数据通过求最后一期，然后求同比
        高频的比率数据，本身就可以边际来看，之前求最后一期
        :return:
        """
        mode = self.mode
        how_map = {'量': 'mean', '价': 'last'}
        freq_pct_map = {'月': 12, '季': 4}
        if mode in ['量', '价']:
            df = self.to_freq(target_freq=freq, how=how_map[mode])
            return df.pct_change(freq_pct_map[freq]).iloc[freq_pct_map[freq]:, :]
        else:
            return self.to_freq(target_freq=freq, how='last')

    def daily_yoy(self, mean_last='mean', mom=False):
        """
        无论原始数据是什么，变频到日度数据后求同比，返回一个元组，
        第一个原始是包含月末时间key的字典，value是DataFrame
        第二个原始是月末值组成的DataFrame
        :return: （dic DataFrame）
        """
        df = self.to_freq('日', mean_last)
        dic = {}
        dic_end = {}
        mld = DateTransformer(df.index[-1]).dt['mld']
        for t in tqdm(pd.date_range(start=df.index[365], end=mld, freq='M')):
            t_dict = {}
            dt = DateTransformer(t).dt
            month_start = dt['mfd']
            month_end = dt['mld']
            if mom:
                se_last_year = df.loc[dt['lmfd']:dt['lmld'], :].mean()  # 上年本月
            else:
                se_last_year = df.loc[dt['lymfd']:dt['lymld'], :].mean()  # 上年本月
            for today in pd.date_range(month_start, month_end, freq='D'):
                se_today = df.loc[month_start:today].mean()
                v = se_today / se_last_year - 1
                t_dict[today] = v
                dic_end[str(t)[:10]] = v
            df_ = pd.DataFrame(t_dict).T
            re_df = df_.fillna(method='ffill').fillna(method='bfill')
            dic[str(t)[:10]] = re_df
        return dic, pd.DataFrame(dic_end).T


if __name__ == "__main__":
    from ini_data import get_data

    data_ = get_data(path='美国通胀.xlsx')
    tran = DataTransformer(data_["美国:核心CPI"])
    data_dic, data_df = tran['月_平均_同比_lag_0']
    dic, df = tran['月_平均_环比_lag_0']
